{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "761a0ad3-7506-4073-bc2f-a46b91452e53",
   "metadata": {},
   "source": [
    "# Instrumentation: Basic Usage"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "21348683-1b4e-42f9-abf8-0548fe02cfde",
   "metadata": {},
   "source": [
    "The `instrumentation` module can be used for observability and monitoring of your llama-index application. It is comprised of the following core abstractions:\n",
    "\n",
    "- `Event` — represents a single moment in time that a certain occurrence took place within the execution of the application’s code.\n",
    "- `EventHandler` — listen to the occurrences of `Event`'s and execute code logic at these moments in time.\n",
    "- `Span` — represents the execution flow of a particular part in the application’s code and thus contains `Event`'s.\n",
    "- `SpanHandler` — is responsible for the entering, exiting, and dropping (i.e., early exiting due to error) of `Span`'s.\n",
    "- `Dispatcher` — emits `Event`'s as well as signals to enter/exit/drop a `Span` to the appropriate handlers."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "0d24c1a2-6dca-4636-b5f6-1547dd045ff3",
   "metadata": {},
   "source": [
    "In this notebook, we demonstrate the basic usage pattern of `instrumentation`:\n",
    "\n",
    "1. Define your custom `EventHandler`\n",
    "2. Define your custom `SpanHandler` which handles an associated `Span` type\n",
    "3. Attach your `EventHandler` and `SpanHandler` to the dispatcher of choice (here, we'll attach it to the root dispatcher)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "9692504e-51a2-4a58-86ae-5de5f5994966",
   "metadata": {},
   "outputs": [],
   "source": [
    "import nest_asyncio\n",
    "\n",
    "nest_asyncio.apply()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "3cb70687-c568-4991-b962-778b5edb3dbf",
   "metadata": {},
   "source": [
    "### Custom Event Handlers"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "8c4ec173-3dc9-49c6-b613-b4a9bf88871f",
   "metadata": {},
   "outputs": [],
   "source": [
    "from llama_index.core.instrumentation.event_handlers import BaseEventHandler"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "7fc325f0-d95c-4d5e-aae6-33d7812ba587",
   "metadata": {},
   "source": [
    "Defining your custom `EventHandler` involves subclassing the `BaseEventHandler`. Doing so, requires defining logic for the abstract method `handle()`."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d5e113aa-f507-4c4e-89a9-bfdc6dbf31ec",
   "metadata": {},
   "outputs": [],
   "source": [
    "class MyEventHandler(BaseEventHandler):\n",
    "    @classmethod\n",
    "    def class_name(cls) -> str:\n",
    "        \"\"\"Class name.\"\"\"\n",
    "        return \"MyEventHandler\"\n",
    "\n",
    "    def handle(self, event) -> None:\n",
    "        \"\"\"Logic for handling event.\"\"\"\n",
    "        # THIS IS WHERE YOU ADD YOUR LOGIC TO HANDLE EVENTS\n",
    "        print(event.dict())\n",
    "        print(\"\")\n",
    "        with open(\"log.txt\", \"a\") as f:\n",
    "            f.write(str(event))\n",
    "            f.write(\"\\n\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "975fb994-436f-4f5d-8a53-74a92ad82a2b",
   "metadata": {},
   "source": [
    "### Custom Span Handlers"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "07799718-c466-4d44-8dfd-b03a24c09eef",
   "metadata": {},
   "source": [
    "`SpanHandler` also involve subclassing a base class, in this case `BaseSpanHandler`. However, since `SpanHandler`'s work with an associated `Span` type, you will need to create this as well if you want to handle a new `Span` type."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d4fc8888-d74c-4bf3-b353-97263104b17e",
   "metadata": {},
   "outputs": [],
   "source": [
    "from llama_index.core.instrumentation.span import BaseSpan"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a22d4b0d-2c64-4aa3-94ea-1e15533be44e",
   "metadata": {},
   "outputs": [],
   "source": [
    "import inspect\n",
    "from typing import Any, Dict, Optional\n",
    "from llama_index.core.bridge.pydantic import Field\n",
    "from llama_index.core.instrumentation.span.base import BaseSpan\n",
    "from llama_index.core.instrumentation.span_handlers import BaseSpanHandler\n",
    "\n",
    "\n",
    "class MyCustomSpan(BaseSpan):\n",
    "    custom_field_1: Any = Field(...)\n",
    "    custom_field_2: Any = Field(...)\n",
    "\n",
    "\n",
    "class MyCustomSpanHandler(BaseSpanHandler[MyCustomSpan]):\n",
    "    @classmethod\n",
    "    def class_name(cls) -> str:\n",
    "        \"\"\"Class name.\"\"\"\n",
    "        return \"MyCustomSpanHandler\"\n",
    "\n",
    "    def new_span(\n",
    "        self,\n",
    "        id_: str,\n",
    "        bound_args: inspect.BoundArguments,\n",
    "        instance: Optional[Any] = None,\n",
    "        parent_span_id: Optional[str] = None,\n",
    "        tags: Optional[Dict[str, Any]] = None,\n",
    "        **kwargs: Any,\n",
    "    ) -> Optional[MyCustomSpan]:\n",
    "        \"\"\"Create a span.\"\"\"\n",
    "        # logic for creating a new MyCustomSpan\n",
    "        pass\n",
    "\n",
    "    def prepare_to_exit_span(\n",
    "        self,\n",
    "        id_: str,\n",
    "        bound_args: inspect.BoundArguments,\n",
    "        instance: Optional[Any] = None,\n",
    "        result: Optional[Any] = None,\n",
    "        **kwargs: Any,\n",
    "    ) -> Any:\n",
    "        \"\"\"Logic for preparing to exit a span.\"\"\"\n",
    "        pass\n",
    "\n",
    "    def prepare_to_drop_span(\n",
    "        self,\n",
    "        id_: str,\n",
    "        bound_args: inspect.BoundArguments,\n",
    "        instance: Optional[Any] = None,\n",
    "        err: Optional[BaseException] = None,\n",
    "        **kwargs: Any,\n",
    "    ) -> Any:\n",
    "        \"\"\"Logic for preparing to drop a span.\"\"\"\n",
    "        pass"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c2e07177-52c1-4d08-8185-4c8e1f6a89fe",
   "metadata": {},
   "source": [
    "For this notebook, we'll use `SimpleSpanHandler` that works with the `SimpleSpan` type."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "18746b28-86e5-4eb8-bebc-8d28435657c9",
   "metadata": {},
   "outputs": [],
   "source": [
    "from llama_index.core.instrumentation.span_handlers import SimpleSpanHandler"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "9ab6d442-6d40-4a19-ac9b-1253ba00028e",
   "metadata": {},
   "source": [
    "### Dispatcher\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "a3710027-218f-4b50-bef6-80497845a376",
   "metadata": {},
   "source": [
    "Now that we have our `EventHandler` and our `SpanHandler`, we can attach it to a `Dispatcher` that will emit `Event`'s and signals to start/exit/drop a `Span` to the appropriate handlers. Those that are familiar with `Logger` from the `logging` Python module, might notice that `Dispatcher` adopts a similar interface. What's more is that `Dispatcher` also utilizes a similar hierarchy and propagation scheme as `Logger`. Specifically, a `dispatcher` will emit `Event`'s to its handlers and by default propagate these events to its parent `Dispatcher` for it to send to its own handlers."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d3138230-2014-4030-9076-59f9d06c0b8d",
   "metadata": {},
   "outputs": [],
   "source": [
    "import llama_index.core.instrumentation as instrument\n",
    "\n",
    "dispatcher = instrument.get_dispatcher()  # modify root dispatcher"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "8d65a394-e390-486b-b261-e2f318c4b9e1",
   "metadata": {},
   "outputs": [],
   "source": [
    "span_handler = SimpleSpanHandler()\n",
    "\n",
    "dispatcher.add_event_handler(MyEventHandler())\n",
    "dispatcher.add_span_handler(span_handler)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "4235053e-4b1e-4280-9284-1a5d9b44cb76",
   "metadata": {},
   "source": [
    "You can also get dispatcher's by name. Purely for the sake of demonstration, in the cells below we get the dispatcher that is defined in the `base.base_query_engine` submodule of `llama_index.core`."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "43798315-04ee-4671-ba43-c02c8d762dbc",
   "metadata": {},
   "outputs": [],
   "source": [
    "qe_dispatcher = instrument.get_dispatcher(\n",
    "    \"llama_index.core.base.base_query_engine\"\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "8c191594-ed8c-4a1e-bdf7-0681bfb76115",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "Dispatcher(name='llama_index.core.base.base_query_engine', event_handlers=[], span_handlers=[], parent_name='root', manager=<llama_index_instrumentation.dispatcher.Manager object at 0x105e5a780>, root_name='root', propagate=True, current_span_ids={})"
      ]
     },
     "execution_count": null,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "qe_dispatcher"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "35d0ba04-2328-41ed-9822-bef7c362fd87",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "Dispatcher(name='root', event_handlers=[NullEventHandler(), MyEventHandler()], span_handlers=[NullSpanHandler(open_spans={}, completed_spans=[], dropped_spans=[], current_span_ids={}), SimpleSpanHandler(open_spans={}, completed_spans=[], dropped_spans=[], current_span_ids={})], parent_name='', manager=None, root_name='root', propagate=False, current_span_ids={})"
      ]
     },
     "execution_count": null,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "qe_dispatcher.parent"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "2ebcd2fa-84d7-459d-9d33-4d02f52c48df",
   "metadata": {},
   "source": [
    "### Test It Out"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a37cb9a7-d5fe-43a1-8da7-4aeae5f74846",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "data/paul_graham/paul_graham_essay.txt: No such file or directory\n"
     ]
    }
   ],
   "source": [
    "!mkdir -p 'data/'\n",
    "!wget 'https://raw.githubusercontent.com/run-llama/llama_index/main/docs/examples/data/paul_graham/paul_graham_essay.txt' -O 'data/paul_graham/paul_graham_essay.txt'"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f02f8dc5-18da-4e70-abed-614e0b98ce8c",
   "metadata": {},
   "outputs": [],
   "source": [
    "from llama_index.core import SimpleDirectoryReader, VectorStoreIndex\n",
    "\n",
    "documents = SimpleDirectoryReader(input_dir=\"./data\").load_data()\n",
    "index = VectorStoreIndex.from_documents(documents)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a3de7479-e48e-42fd-8515-f04e1c314469",
   "metadata": {},
   "outputs": [],
   "source": [
    "query_engine = index.as_query_engine()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "e44df059-edc6-4342-ac06-8a9f293b8d5d",
   "metadata": {},
   "source": [
    "#### Sync"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "5719b4d5-d50b-49bc-80ca-a1233817c64d",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "{'timestamp': datetime.datetime(2024, 3, 14, 15, 57, 50, 614289), 'id_': UUID('35643a53-52da-4547-b770-dba7600c9070'), 'class_name': 'QueryStartEvent'}\n",
      "\n",
      "{'timestamp': datetime.datetime(2024, 3, 14, 15, 57, 50, 615966), 'id_': UUID('18fa9b70-6fbc-4ce7-9e45-1f292766e820'), 'class_name': 'RetrievalStartEvent'}\n",
      "\n",
      "{'timestamp': datetime.datetime(2024, 3, 14, 15, 57, 50, 810658), 'id_': UUID('2aed6810-99a1-49ab-a6c4-5f242f1b1de3'), 'class_name': 'RetrievalEndEvent'}\n",
      "\n",
      "{'timestamp': datetime.datetime(2024, 3, 14, 15, 57, 50, 811651), 'id_': UUID('6f8d96a2-4da0-485f-9a73-4dae2d026b48'), 'class_name': 'SynthesizeStartEvent'}\n",
      "\n",
      "{'timestamp': datetime.datetime(2024, 3, 14, 15, 57, 50, 818960), 'id_': UUID('005e52bd-25f6-49cb-8766-5399098b7e51'), 'class_name': 'GetResponseStartEvent'}\n",
      "\n",
      "{'timestamp': datetime.datetime(2024, 3, 14, 15, 57, 50, 823964), 'id_': UUID('52273f18-2865-42d3-a11e-acbe5bb56748'), 'class_name': 'LLMPredictStartEvent'}\n",
      "\n",
      "{'timestamp': datetime.datetime(2024, 3, 14, 15, 57, 52, 375382), 'id_': UUID('54516719-425b-422a-b38c-8bd002d8fa74'), 'class_name': 'LLMPredictEndEvent'}\n",
      "\n",
      "{'timestamp': datetime.datetime(2024, 3, 14, 15, 57, 52, 376003), 'id_': UUID('24ec9eb3-d2ad-46f7-8db2-faf68075c7b3'), 'class_name': 'GetResponseEndEvent'}\n",
      "\n",
      "{'timestamp': datetime.datetime(2024, 3, 14, 15, 57, 52, 376347), 'id_': UUID('6d453292-374c-4fef-8bc4-ae164455a133'), 'class_name': 'SynthesizeEndEvent'}\n",
      "\n",
      "{'timestamp': datetime.datetime(2024, 3, 14, 15, 57, 52, 376505), 'id_': UUID('9919903a-dbb1-4bdf-8029-3e4fcd9ea073'), 'class_name': 'QueryEndEvent'}\n",
      "\n"
     ]
    }
   ],
   "source": [
    "query_result = query_engine.query(\"Who is Paul?\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "3810b999-d43e-4ca7-970d-4ccc41f84027",
   "metadata": {},
   "source": [
    "#### Async"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "5f59bdfe-c661-4547-ae42-f8143a4b1c2c",
   "metadata": {},
   "source": [
    "`Dispatcher` also works on async methods."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "23e70e94-e116-4e00-b75d-d13a3b1c23d8",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "{'timestamp': datetime.datetime(2024, 3, 14, 15, 58, 35, 276918), 'id_': UUID('8ea98ded-91f2-45cf-b418-b92b3c7693bf'), 'class_name': 'QueryStartEvent'}\n",
      "\n",
      "{'timestamp': datetime.datetime(2024, 3, 14, 15, 58, 35, 279006), 'id_': UUID('60ff04ce-0972-4fe1-bfaf-5ffaaea3ef4a'), 'class_name': 'RetrievalStartEvent'}\n",
      "\n",
      "{'timestamp': datetime.datetime(2024, 3, 14, 15, 58, 35, 555879), 'id_': UUID('c29bb55e-b2e1-4637-a6cb-745a2424da90'), 'class_name': 'RetrievalEndEvent'}\n",
      "\n",
      "{'timestamp': datetime.datetime(2024, 3, 14, 15, 58, 35, 557244), 'id_': UUID('dee202a1-f760-495e-a6f8-3644a535ff13'), 'class_name': 'SynthesizeStartEvent'}\n",
      "\n",
      "{'timestamp': datetime.datetime(2024, 3, 14, 15, 58, 35, 564098), 'id_': UUID('c524828c-cdd7-4ddd-876a-5fda5f10143d'), 'class_name': 'GetResponseStartEvent'}\n",
      "\n",
      "{'timestamp': datetime.datetime(2024, 3, 14, 15, 58, 35, 568930), 'id_': UUID('bd8e78ce-9a87-41a8-b009-e0694e09e0b3'), 'class_name': 'LLMPredictStartEvent'}\n",
      "\n",
      "{'timestamp': datetime.datetime(2024, 3, 14, 15, 58, 37, 70264), 'id_': UUID('54633e8f-8f89-4a4c-9e79-f059f9dbecc2'), 'class_name': 'LLMPredictEndEvent'}\n",
      "\n",
      "{'timestamp': datetime.datetime(2024, 3, 14, 15, 58, 37, 71236), 'id_': UUID('0a110c5c-3d4c-4eeb-8066-b8e512e69838'), 'class_name': 'GetResponseEndEvent'}\n",
      "\n",
      "{'timestamp': datetime.datetime(2024, 3, 14, 15, 58, 37, 71652), 'id_': UUID('1152f15e-ac5b-4292-ad9f-45f8404183f9'), 'class_name': 'SynthesizeEndEvent'}\n",
      "\n",
      "{'timestamp': datetime.datetime(2024, 3, 14, 15, 58, 37, 71891), 'id_': UUID('8aa8f930-1ac2-4924-a185-00a06bc7ba79'), 'class_name': 'QueryEndEvent'}\n",
      "\n"
     ]
    }
   ],
   "source": [
    "query_result = await query_engine.aquery(\"Who is Paul?\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "f2fcb260-cec3-446a-9971-49bad189e83d",
   "metadata": {},
   "source": [
    "#### Streaming"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "0b11b898-9a55-4ac4-b7ed-aa92ca1b6865",
   "metadata": {},
   "source": [
    "`Dispatcher` also works on methods that support streaming!"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "fef4f47b-291a-4878-81e0-bef7cb5299e4",
   "metadata": {},
   "outputs": [],
   "source": [
    "chat_engine = index.as_chat_engine()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "491e0962-d49d-428f-b861-b9883e1708fc",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "{'timestamp': datetime.datetime(2024, 3, 14, 15, 59, 31, 345865), 'id_': UUID('1d9643ca-368b-4e08-9878-ed7682196007'), 'class_name': 'AgentChatWithStepStartEvent'}\n",
      "\n",
      "{'timestamp': datetime.datetime(2024, 3, 14, 15, 59, 31, 346727), 'id_': UUID('c38a18e3-0c2c-43b3-a1a9-0fb2e696e627'), 'class_name': 'AgentRunStepStartEvent'}\n",
      "\n",
      "{'timestamp': datetime.datetime(2024, 3, 14, 15, 59, 31, 348524), 'id_': UUID('9a49dd15-715c-474a-b704-b9e8df919c50'), 'class_name': 'StreamChatStartEvent'}\n",
      "\n",
      "{'timestamp': datetime.datetime(2024, 3, 14, 15, 59, 31, 975148), 'id_': UUID('c3b5b9eb-104d-461e-a081-e65ec9dc3131'), 'class_name': 'StreamChatEndEvent'}\n",
      "\n",
      "{'timestamp': datetime.datetime(2024, 3, 14, 15, 59, 31, 977522), 'id_': UUID('2e6c3935-8ece-49bf-aacc-de18fd79da42'), 'class_name': 'QueryStartEvent'}\n",
      "\n",
      "{'timestamp': datetime.datetime(2024, 3, 14, 15, 59, 31, 978389), 'id_': UUID('5c43f441-d262-427f-8ef7-b3b6e6e86f44'), 'class_name': 'RetrievalStartEvent'}\n",
      "\n",
      "{'timestamp': datetime.datetime(2024, 3, 14, 15, 59, 32, 188462), 'id_': UUID('724dd93f-39a8-4b12-b056-194c4cf3ed72'), 'class_name': 'RetrievalEndEvent'}\n",
      "\n",
      "{'timestamp': datetime.datetime(2024, 3, 14, 15, 59, 32, 189601), 'id_': UUID('27e5ac36-d313-4df6-bc8b-40b79e59698e'), 'class_name': 'SynthesizeStartEvent'}\n",
      "\n",
      "{'timestamp': datetime.datetime(2024, 3, 14, 15, 59, 32, 208520), 'id_': UUID('77ec49c0-fb24-46dd-b843-954e241a0eb0'), 'class_name': 'GetResponseStartEvent'}\n",
      "\n",
      "{'timestamp': datetime.datetime(2024, 3, 14, 15, 59, 32, 214106), 'id_': UUID('b28106fa-8d8e-49ca-92af-a37e0db45b9f'), 'class_name': 'LLMPredictStartEvent'}\n",
      "\n",
      "{'timestamp': datetime.datetime(2024, 3, 14, 15, 59, 33, 59544), 'id_': UUID('0e40cac7-9eb4-48d3-81bc-d01a5b3c0440'), 'class_name': 'LLMPredictEndEvent'}\n",
      "\n",
      "{'timestamp': datetime.datetime(2024, 3, 14, 15, 59, 33, 60941), 'id_': UUID('77e1d7eb-5807-4184-9e18-de4fdde18654'), 'class_name': 'GetResponseEndEvent'}\n",
      "\n",
      "{'timestamp': datetime.datetime(2024, 3, 14, 15, 59, 33, 61349), 'id_': UUID('565c9019-89b0-4dec-bb54-0d9642030009'), 'class_name': 'SynthesizeEndEvent'}\n",
      "\n",
      "{'timestamp': datetime.datetime(2024, 3, 14, 15, 59, 33, 61677), 'id_': UUID('3cba488e-1d69-4b75-a601-2cf467816ef0'), 'class_name': 'QueryEndEvent'}\n",
      "\n",
      "{'timestamp': datetime.datetime(2024, 3, 14, 15, 59, 33, 62157), 'id_': UUID('a550a4c7-40b3-43ac-abb5-5cab4b759888'), 'class_name': 'AgentRunStepEndEvent'}\n",
      "\n",
      "{'timestamp': datetime.datetime(2024, 3, 14, 15, 59, 33, 62417), 'id_': UUID('d99b6fd9-bf35-4c37-9625-da39cc5aef23'), 'class_name': 'AgentRunStepStartEvent'}\n",
      "\n",
      "{'timestamp': datetime.datetime(2024, 3, 14, 15, 59, 33, 63294), 'id_': UUID('b8015c1d-53b7-4530-812c-2ded8dab4c67'), 'class_name': 'StreamChatStartEvent'}\n",
      "\n",
      "{'timestamp': datetime.datetime(2024, 3, 14, 15, 59, 33, 387260), 'id_': UUID('d5bbdb58-be35-444d-b0f0-0ccbc570e54e'), 'delta': 'Why', 'class_name': 'StreamChatDeltaReceivedEvent'}\n",
      "\n",
      "{'timestamp': datetime.datetime(2024, 3, 14, 15, 59, 33, 389911), 'id_': UUID('91e78251-def6-4bfb-9712-20c4c0d2690a'), 'class_name': 'AgentRunStepEndEvent'}\n",
      "\n",
      "{'timestamp': datetime.datetime(2024, 3, 14, 15, 59, 33, 389700), 'id_': UUID('8b181253-1ea6-4ffb-b384-c425307d7b88'), 'delta': ' did', 'class_name': 'StreamChatDeltaReceivedEvent'}\n",
      "\n",
      "{'timestamp': datetime.datetime(2024, 3, 14, 15, 59, 33, 390495), 'id_': UUID('57dd72cf-8b8c-4596-8de7-4a701fc50125'), 'class_name': 'AgentChatWithStepEndEvent'}\n",
      "\n",
      "{'timestamp': datetime.datetime(2024, 3, 14, 15, 59, 33, 409497), 'id_': UUID('e6ceba61-5c78-4ee3-972e-c9e02dbddc7c'), 'delta': ' the', 'class_name': 'StreamChatDeltaReceivedEvent'}\n",
      "\n",
      "{'timestamp': datetime.datetime(2024, 3, 14, 15, 59, 33, 410653), 'id_': UUID('e9d8d9fe-1080-455c-8add-06b1774506bc'), 'delta': ' computer', 'class_name': 'StreamChatDeltaReceivedEvent'}\n",
      "\n",
      "{'timestamp': datetime.datetime(2024, 3, 14, 15, 59, 33, 449414), 'id_': UUID('04cffd87-ca8a-4efb-af2a-a99964610e4c'), 'delta': ' keep', 'class_name': 'StreamChatDeltaReceivedEvent'}\n",
      "\n",
      "{'timestamp': datetime.datetime(2024, 3, 14, 15, 59, 33, 450316), 'id_': UUID('21824a0a-3674-42d1-91c7-d8a865c2b270'), 'delta': ' its', 'class_name': 'StreamChatDeltaReceivedEvent'}\n",
      "\n",
      "{'timestamp': datetime.datetime(2024, 3, 14, 15, 59, 33, 495431), 'id_': UUID('a677507c-6b74-454f-a185-3df008e9e5ff'), 'delta': ' drinks', 'class_name': 'StreamChatDeltaReceivedEvent'}\n",
      "\n",
      "{'timestamp': datetime.datetime(2024, 3, 14, 15, 59, 33, 496188), 'id_': UUID('8b94885c-ce78-46cc-8ee6-938484ae2980'), 'delta': ' on', 'class_name': 'StreamChatDeltaReceivedEvent'}\n",
      "\n",
      "{'timestamp': datetime.datetime(2024, 3, 14, 15, 59, 33, 527857), 'id_': UUID('dc356d5d-c968-4d43-9903-bc158b73338a'), 'delta': ' the', 'class_name': 'StreamChatDeltaReceivedEvent'}\n",
      "\n",
      "{'timestamp': datetime.datetime(2024, 3, 14, 15, 59, 33, 529075), 'id_': UUID('f8a2637f-e746-418d-8921-4e7e8c26956f'), 'delta': ' motherboard', 'class_name': 'StreamChatDeltaReceivedEvent'}\n",
      "\n",
      "{'timestamp': datetime.datetime(2024, 3, 14, 15, 59, 33, 554042), 'id_': UUID('dd21cd77-c329-4947-ad77-7683885e0ce1'), 'delta': '?', 'class_name': 'StreamChatDeltaReceivedEvent'}\n",
      "\n",
      "{'timestamp': datetime.datetime(2024, 3, 14, 15, 59, 33, 557320), 'id_': UUID('93a2a7e1-7fe5-4eb8-851e-1b3dea4df6b5'), 'delta': ' Because', 'class_name': 'StreamChatDeltaReceivedEvent'}\n",
      "\n",
      "{'timestamp': datetime.datetime(2024, 3, 14, 15, 59, 33, 608305), 'id_': UUID('445cb52e-9cf3-4f85-aba3-6383974e6b5f'), 'delta': ' it', 'class_name': 'StreamChatDeltaReceivedEvent'}\n",
      "\n",
      "{'timestamp': datetime.datetime(2024, 3, 14, 15, 59, 33, 609392), 'id_': UUID('bd7bb3dc-5bd7-418e-a18a-203cdb701bcb'), 'delta': ' had', 'class_name': 'StreamChatDeltaReceivedEvent'}\n",
      "\n",
      "{'timestamp': datetime.datetime(2024, 3, 14, 15, 59, 33, 609896), 'id_': UUID('e981bdb3-57ca-456e-8353-be087ea6bb55'), 'delta': ' too', 'class_name': 'StreamChatDeltaReceivedEvent'}\n",
      "\n",
      "{'timestamp': datetime.datetime(2024, 3, 14, 15, 59, 33, 610659), 'id_': UUID('79b3eb21-6f8d-4a87-9f26-04e757a29da3'), 'delta': ' many', 'class_name': 'StreamChatDeltaReceivedEvent'}\n",
      "\n",
      "{'timestamp': datetime.datetime(2024, 3, 14, 15, 59, 33, 667840), 'id_': UUID('73073cd6-dfae-4632-a302-a4841a08f272'), 'delta': ' bytes', 'class_name': 'StreamChatDeltaReceivedEvent'}\n",
      "\n",
      "{'timestamp': datetime.datetime(2024, 3, 14, 15, 59, 33, 669579), 'id_': UUID('c30da40c-bf7d-49f0-9505-2345fc67fde7'), 'delta': '!', 'class_name': 'StreamChatDeltaReceivedEvent'}\n",
      "\n",
      "{'timestamp': datetime.datetime(2024, 3, 14, 15, 59, 33, 670733), 'id_': UUID('f9baf8fc-8755-4740-8fe4-f5a3c9d77f9e'), 'delta': ' 😄', 'class_name': 'StreamChatDeltaReceivedEvent'}\n",
      "\n",
      "{'timestamp': datetime.datetime(2024, 3, 14, 15, 59, 33, 672180), 'id_': UUID('e07eb812-68b5-4027-8ee0-87bf6cd5c744'), 'class_name': 'StreamChatEndEvent'}\n",
      "\n"
     ]
    }
   ],
   "source": [
    "streaming_response = chat_engine.stream_chat(\"Tell me a joke.\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "27e4c33f-ee90-4716-8dc9-2ed22edd5197",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Why did the computer keep its drinks on the motherboard? Because it had too many bytes! 😄"
     ]
    }
   ],
   "source": [
    "for token in streaming_response.response_gen:\n",
    "    print(token, end=\"\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "d2266149-be2c-462b-a029-652363cb925f",
   "metadata": {},
   "source": [
    "### Printing Basic Trace Trees with `SimpleSpanHandler`"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "876ee4dd-3e28-4aae-a917-daac42f17b6f",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "BaseQueryEngine.query-bda10f51-e5c8-4ef8-9467-c816b0c92797 (1.762367)\n",
      "└── RetrieverQueryEngine._query-35da82df-8e64-43c5-9046-11df14b3b9f7 (1.760649)\n",
      "    ├── BaseRetriever.retrieve-237d1b8d-f8b1-4e0b-908c-00087f086c2c (0.19558)\n",
      "    │   └── VectorIndexRetriever._retrieve-af6479b8-3210-41df-9d6f-3af31059f274 (0.194024)\n",
      "    └── BaseSynthesizer.synthesize-bf923672-6e60-4015-b6fe-0d0c9d1d35e3 (1.564853)\n",
      "        └── CompactAndRefine.get_response-447c173e-4016-4376-9c56-a7171ea1ddf0 (1.564162)\n",
      "            └── Refine.get_response-83b1159d-d33f-401f-a76e-bc3ee5095b57 (1.557365)\n",
      "                └── LLM.predict-a5ab2252-1eb1-4413-9ef0-efa14c2c3b6b (1.552019)\n",
      "\n",
      "\n",
      "BaseQueryEngine.aquery-7f3daee7-540f-4189-a350-a37e7e0596d5 (1.79559)\n",
      "└── RetrieverQueryEngine._aquery-cc049d88-933a-41d3-abb4-06c5bd567e45 (1.793149)\n",
      "    ├── BaseRetriever.aretrieve-c1a2ae34-3916-4069-81ba-7ba9b9d7235d (0.278098)\n",
      "    │   └── VectorIndexRetriever._aretrieve-ef38d533-dd12-4fa5-9879-cd885124d8aa (0.276331)\n",
      "    └── BaseSynthesizer.asynthesize-e3f02693-1563-4eec-898e-1043bbeec870 (1.514635)\n",
      "        └── CompactAndRefine.aget_response-6cfcc5f8-1a47-4dde-aca7-8bd6543ce457 (1.513896)\n",
      "            └── Refine.aget_response-7c2a4f67-f4bb-4065-b934-ac8dae7a9529 (1.507486)\n",
      "                └── LLM.apredict-3ccf1ad0-d0b3-44bd-b64a-83bb652ed4c8 (1.502215)\n",
      "\n",
      "\n",
      "AgentRunner.stream_chat-ae3bc1f0-b9ff-456c-a3b4-7264a5757336 (2.045523)\n",
      "└── AgentRunner._chat-dd1d6afa-276a-4f7b-8c01-eb39df07b74d (2.045444)\n",
      "    ├── AgentRunner._run_step-be8844fc-bd2c-4845-b5c5-378e9a1c97b5 (1.715629)\n",
      "    │   ├── StreamingAgentChatResponse.write_response_to_history-b3fddc19-ee88-442b-9c95-0ea7e68fe4f3 (0.62843)\n",
      "    │   └── BaseQueryEngine.query-4eb2f0ea-89c0-48d8-a757-92961a4c4275 (1.084424)\n",
      "    │       └── RetrieverQueryEngine._query-fc5f2290-b715-4d33-82a6-d0eea3ba8625 (1.083421)\n",
      "    │           ├── BaseRetriever.retrieve-62017ac0-3d6b-4ca4-91e9-d0548d226536 (0.211132)\n",
      "    │           │   └── VectorIndexRetriever._retrieve-c77b4ae8-702c-42a9-b50a-a394c031c728 (0.209355)\n",
      "    │           └── BaseSynthesizer.synthesize-7b6caba0-8156-4611-9f8a-baa3a7e5e151 (0.872036)\n",
      "    │               └── CompactAndRefine.get_response-ec9651ef-4a88-4395-89da-52c2336baca3 (0.871197)\n",
      "    │                   └── Refine.get_response-57a69243-b676-4f33-8c61-73a5ba6af03c (0.852795)\n",
      "    │                       └── LLM.predict-b0fe909f-31e5-4655-8428-5dee6f6de7c8 (0.847305)\n",
      "    └── AgentRunner._run_step-9a126f49-c704-45a8-9104-0e0a9d19c956 (0.327996)\n",
      "        └── StreamingAgentChatResponse.write_response_to_history-b727aeea-dd0a-4c7a-9212-cf591caf7bb5 (0.609362)\n",
      "\n",
      "\n"
     ]
    }
   ],
   "source": [
    "span_handler.print_trace_trees()"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": ".venv",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
